#!/usr/bin/env python
# -*- coding: utf-8 -*-

from typing import Dict, Any, List, Optional, AsyncGenerator
import asyncio
from functools import wraps

from pymilvus import connections, Collection, utility
from pymilvus.exceptions import MilvusException

from app.core.config import settings
from app.core.exceptions import VectorDBError
from app.core.logging import get_logger

# 创建日志记录器
logger = get_logger("milvus")


class MilvusClient:
    """Milvus客户端封装类"""
    
    def __init__(self):
        """初始化Milvus客户端"""
        self.host = settings.MILVUS_HOST
        self.port = settings.MILVUS_PORT
        self.user = settings.MILVUS_USER
        self.password = settings.MILVUS_PASSWORD
        self.connected = False
        self.connection_alias = "default"
        
        # 连接参数
        self.connect_params = {
            "host": self.host,
            "port": self.port,
        }
        
        # 如果提供了用户名和密码，则添加到连接参数中
        if self.user and self.password:
            self.connect_params["user"] = self.user
            self.connect_params["password"] = self.password
    
    async def connect(self) -> None:
        """连接到Milvus服务器"""
        if self.connected:
            return
        
        try:
            # 使用线程池执行同步操作
            loop = asyncio.get_event_loop()
            await loop.run_in_executor(
                None,
                lambda: connections.connect(
                    alias=self.connection_alias,
                    **self.connect_params
                )
            )
            self.connected = True
            logger.info(f"已连接到Milvus服务器: {self.host}:{self.port}")
        except Exception as e:
            logger.error(f"连接Milvus服务器失败: {str(e)}")
            raise VectorDBError(f"连接Milvus服务器失败: {str(e)}")
    
    async def disconnect(self) -> None:
        """断开与Milvus服务器的连接"""
        if not self.connected:
            return
        
        try:
            # 使用线程池执行同步操作
            loop = asyncio.get_event_loop()
            await loop.run_in_executor(
                None,
                lambda: connections.disconnect(self.connection_alias)
            )
            self.connected = False
            logger.info("已断开与Milvus服务器的连接")
        except Exception as e:
            logger.error(f"断开Milvus服务器连接失败: {str(e)}")
    
    async def ensure_connected(self) -> None:
        """确保已连接到Milvus服务器"""
        if not self.connected:
            await self.connect()
    
    async def collection_exists(self, collection_name: str) -> bool:
        """检查集合是否存在
        
        Args:
            collection_name: 集合名称
            
        Returns:
            bool: 集合是否存在
        """
        await self.ensure_connected()
        
        try:
            # 使用线程池执行同步操作
            loop = asyncio.get_event_loop()
            exists = await loop.run_in_executor(
                None,
                lambda: utility.has_collection(collection_name)
            )
            return exists
        except Exception as e:
            logger.error(f"检查集合是否存在失败: {str(e)}")
            raise VectorDBError(f"检查集合是否存在失败: {str(e)}")
    
    async def get_collection(self, collection_name: str) -> Collection:
        """获取集合
        
        Args:
            collection_name: 集合名称
            
        Returns:
            Collection: Milvus集合对象
            
        Raises:
            VectorDBError: 如果集合不存在或获取失败
        """
        await self.ensure_connected()
        
        # 检查集合是否存在
        exists = await self.collection_exists(collection_name)
        if not exists:
            raise VectorDBError(f"集合不存在: {collection_name}")
        
        try:
            # 使用线程池执行同步操作
            loop = asyncio.get_event_loop()
            collection = await loop.run_in_executor(
                None,
                lambda: Collection(collection_name)
            )
            
            # 加载集合
            await loop.run_in_executor(
                None,
                lambda: collection.load()
            )
            
            return collection
        except Exception as e:
            logger.error(f"获取集合失败: {str(e)}")
            raise VectorDBError(f"获取集合失败: {str(e)}")
    
    async def create_collection(
        self,
        collection_name: str,
        dimension: int,
        description: str = ""
    ) -> Collection:
        """创建集合
        
        Args:
            collection_name: 集合名称
            dimension: 向量维度
            description: 集合描述
            
        Returns:
            Collection: 创建的集合对象
            
        Raises:
            VectorDBError: 如果创建失败
        """
        await self.ensure_connected()
        
        # 检查集合是否已存在
        exists = await self.collection_exists(collection_name)
        if exists:
            logger.info(f"集合已存在: {collection_name}")
            return await self.get_collection(collection_name)
        
        try:
            from pymilvus import CollectionSchema, FieldSchema, DataType
            
            # 定义字段
            fields = [
                FieldSchema(name="id", dtype=DataType.INT64, is_primary=True, auto_id=True),
                FieldSchema(name="content_id", dtype=DataType.INT64, description="内容ID"),
                FieldSchema(name="data_type", dtype=DataType.VARCHAR, max_length=20, description="数据类型"),
                FieldSchema(name="embedding", dtype=DataType.FLOAT_VECTOR, dim=dimension, description="向量嵌入")
            ]
            
            # 创建集合模式
            schema = CollectionSchema(fields=fields, description=description)
            
            # 使用线程池执行同步操作
            loop = asyncio.get_event_loop()
            collection = await loop.run_in_executor(
                None,
                lambda: Collection(name=collection_name, schema=schema)
            )
            
            # 创建索引
            index_params = {
                "metric_type": "COSINE",
                "index_type": "HNSW",
                "params": {"M": 8, "efConstruction": 64}
            }
            
            await loop.run_in_executor(
                None,
                lambda: collection.create_index(field_name="embedding", index_params=index_params)
            )
            
            # 加载集合
            await loop.run_in_executor(
                None,
                lambda: collection.load()
            )
            
            logger.info(f"已创建集合: {collection_name}")
            return collection
        except Exception as e:
            logger.error(f"创建集合失败: {str(e)}")
            raise VectorDBError(f"创建集合失败: {str(e)}")
    
    async def drop_collection(self, collection_name: str) -> None:
        """删除集合
        
        Args:
            collection_name: 集合名称
            
        Raises:
            VectorDBError: 如果删除失败
        """
        await self.ensure_connected()
        
        # 检查集合是否存在
        exists = await self.collection_exists(collection_name)
        if not exists:
            logger.info(f"集合不存在，无需删除: {collection_name}")
            return
        
        try:
            # 使用线程池执行同步操作
            loop = asyncio.get_event_loop()
            await loop.run_in_executor(
                None,
                lambda: utility.drop_collection(collection_name)
            )
            logger.info(f"已删除集合: {collection_name}")
        except Exception as e:
            logger.error(f"删除集合失败: {str(e)}")
            raise VectorDBError(f"删除集合失败: {str(e)}")


# 全局Milvus客户端实例
milvus_client = MilvusClient()


async def get_milvus_client() -> MilvusClient:
    """获取Milvus客户端
    
    Returns:
        MilvusClient: Milvus客户端实例
    """
    await milvus_client.ensure_connected()
    return milvus_client